import cv2
import numpy as np
from tqdm import tqdm
import tensorflow as tf  # 导入TensorFlow库

# 读取视频
cap = cv2.VideoCapture('Y:/MP4/1.mp4')
fps = int(cap.get(cv2.CAP_PROP_FPS))
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print(f'frame_width:{frame_width},frame_height:{frame_height},fps:{fps}')

# 输出视频编码器设置
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter('Y:/MP4/1_no_watermark.mp4', fourcc, fps, (frame_width, frame_height))

# 使用tqdm显示进度条
with tqdm(total=total_frames, desc="Processing Video") as pbar:
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        # 假设我们知道水印的位置和大小
        y1, y2, x1, x2 = 411, 481, 652, 909  # 修正后的坐标
        # 对每一帧进行去水印处理
        mask = np.zeros(frame.shape[:2], np.uint8)
        cv2.rectangle(mask, (x1, y1), (x2, y2), 255, -1)  # 水印区域掩码
        # 使用TensorFlow进行去水印
        frame_tensor = tf.convert_to_tensor(frame)
        mask_tensor = tf.convert_to_tensor(mask)
        # 使用TensorFlow的图像处理功能
        # 使用OpenCV的inpaint方法
        dst = cv2.inpaint(frame, mask, 3, cv2.INPAINT_TELEA)
        out.write(dst)
        pbar.update(1)  # 更新进度条

cap.release()
out.release()